package cn.qyd.leecode;

import jdk.internal.dynalink.linker.LinkerServices;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/**
 * @author 邱运铎
 * @date 2024-04-17 20:03
 */
public class BadNodeInfectII01 {
    public static void main(String[] args) {
        BadNodeInfectII01 demo = new BadNodeInfectII01();
//        int[][] graph = new int[][]{{1,1,0},{1,1,0},{0,0,1}};
//        int[] initial = new int[]{0, 1};
        int[][] graph = new int[][]{{1,1,0,0,0,0,0,0,0},
                {1,1,0,0,0,0,0,0,0},
                {0,0,1,1,0,0,0,0,0},
                {0,0,1,1,1,0,0,0,0},
                {0,0,0,1,1,0,0,0,0},
                {0,0,0,0,1,1,0,0,0},
                {0,0,0,0,0,0,1,1,0},
                {0,0,0,0,0,0,1,1,1},
                {0,0,0,0,0,0,0,1,1}};
        int[] initial = new int[]{0,1,3,5};
        int i = demo.minMalwareSpread(graph, initial);
        System.out.println(i);
    }

    public int minMalwareSpread(int[][] graph, int[] initial) {
        int n = graph.length;
        boolean[] isInitial = new boolean[n];
        for (int i : initial) {
            isInitial[i] = true;
        }

        int[] uf = new int[n];
        for (int i = 0; i < n; i++) {
            uf[i] = i;
        }

        for (int i = 0; i < n; i++) {
            if (isInitial[i]) {
                continue;
            }
            for (int j = 0; j < n; j++) {
                if (isInitial[j]) {
                    continue;
                }
                if (graph[i][j] == 1) {
                    merge(uf, i, j);
                }
            }
        }
        System.out.println(Arrays.toString(uf));

        List<Integer>[] infectedBy = new List[n];
        Arrays.setAll(infectedBy, l -> new ArrayList<>());

        for (int initialBadNode : initial) {
            boolean[] infectedSet = new boolean[n];
            for (int i = 0; i < n; i++) {
                if (isInitial[i] || graph[i][initialBadNode] == 0) {
                    continue;
                }
                infectedSet[find(uf, i)] = true;
            }
            for (int i = 0; i < n; i++) {
                if (infectedSet[i]) {
                    infectedBy[i].add(initialBadNode);
                }
            }
        }
        System.out.println(Arrays.toString(infectedBy));
        int[] count = new int[n];
        for (int i = 0; i < n; i++) {
            if (infectedBy[i].size() != 1) {
                continue;
            }
            int firstNode = infectedBy[i].get(0);
            for (int j = 0; j < n; j++) {
                if (find(uf, j) == find(uf, i)) {
                    count[firstNode]++;
                }
            }
        }
        System.out.println(Arrays.toString(count));

        int removeNode = initial[0];
        for (int initialBadNode : initial) {
            if (count[initialBadNode] > count[removeNode] || count[initialBadNode] == count[removeNode] && initialBadNode < removeNode) {
                removeNode = initialBadNode;
            }
        }

        return removeNode;

    }

    public int find(int[] uf, int i) {
        if (uf[i] == i) {
            return i;
        }

        uf[i] = find(uf, uf[i]);
        return uf[i];
    }

    public void merge(int[] uf, int x, int y) {
        int rx = find(uf, x);
        int ry = find(uf, y);
        uf[rx] = ry;
    }
}
